
#-*-coding:utf-8-*-#

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import container
plt.rcParams['font.size'] = 14

os.chdir('PATH')

"""

Drought simulation results: Data_Fig3ab

Wildfire simulation results: Data_Fig3cd

"""


dataframes = {}

for i in range(2, 16):
    dataframes[f'dr{i}'] = pd.read_csv('drought simulation AIDD '+str(i)+'.csv',index_col=0)

for i in range(2, 16):
    dataframes[f'wf{i}'] = pd.read_csv('wildfire simulation CAIDD '+str(i)+'.csv',index_col=0)

gammadataframes = {}

for i in range(2, 16):
    gammadataframes[f'dr{i}'] = pd.read_csv('drought simulation AIDD b '+str(i)+'.csv',index_col=0)

for i in range(2, 16):
    gammadataframes[f'wf{i}'] = pd.read_csv('wildfire simulation AIDD b '+str(i)+'.csv',index_col=0)


taudataframes = {}

for i in range(2, 16):
    taudataframes[f'dr{i}'] = pd.read_csv('drought simulation AIDD t '+str(i)+'.csv',index_col=0)

for i in range(2, 16):
    taudataframes[f'wf{i}'] = pd.read_csv('wildfire simulation AIDD t '+str(i)+'.csv',index_col=0)



def sens_test(df,simulation_num,leng):
    sdf = pd.DataFrame()
    sens = 0
    sign = 0
    for i in range(leng):
        if sens == 0 and 0 < df[str(simulation_num) + 't'].iloc[i]*100 and 0 < df[str(simulation_num) + 'b'].iloc[i]*100:
            sens = 1
        if sens == 1 and 0 < df[str(simulation_num) + 't'].iloc[i]*100 and 0 > df[str(simulation_num) + 'b'].iloc[i]*100:
            sign = i
            break
    if sens == 1:
        sdf.loc['vt','value'] = df[str(simulation_num) + 't'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['vm','value'] = df[str(simulation_num) + 'm'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['vb','value'] = df[str(simulation_num) + 'b'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['d','value'] = df[str(simulation_num) + 'm'].idxmax()
        if sign == 0:
            sdf.loc['sign','value'] = 194
        else:
            sdf.loc['sign','value'] = sign
        return sdf
    else:
        sdf.loc['vt','value'] = df[str(simulation_num) + 't'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['vm','value'] = df[str(simulation_num) + 'm'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['vb','value'] = df[str(simulation_num) + 'b'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['d','value'] = df[str(simulation_num) + 'm'].idxmax()
        sdf.loc['sign','value'] = -1
        return sdf

sendf = {}

for i in range(2,16):
    sendf[f'drsdf{i}'] = sens_test(dataframes['dr'+str(i)], 3, len(dataframes['dr'+str(i)]))

for i in range(2,16):
    sendf[f'wdrsdf{i}'] = sens_test(dataframes['dr'+str(i)], 2, len(dataframes['dr'+str(i)]))


drsdf = pd.concat([sendf['drsdf2'],sendf['drsdf3'],sendf['drsdf4'],sendf['drsdf5'],sendf['drsdf6'],sendf['drsdf7'],
                   sendf['drsdf8'],sendf['drsdf9'],sendf['drsdf10'],sendf['drsdf11'],sendf['drsdf12'],
                   sendf['drsdf13'],sendf['drsdf14'],sendf['drsdf15']],axis=1)
drsdf.columns = range(2,16,1)

wdrsdf = pd.concat([sendf['wdrsdf2'],sendf['wdrsdf3'],sendf['wdrsdf4'],sendf['wdrsdf5'],sendf['wdrsdf6'],sendf['wdrsdf7'],
                   sendf['wdrsdf8'],sendf['wdrsdf9'],sendf['wdrsdf10'],sendf['wdrsdf11'],sendf['wdrsdf12'],
                   sendf['wdrsdf13'],sendf['wdrsdf14'],sendf['wdrsdf15']],axis=1)
wdrsdf.columns = range(2,16,1)

drsdf = drsdf.T
wdrsdf = wdrsdf.T


for i in range(2,16):
    sendf[f'wfsdf{i}'] = sens_test(dataframes['wf'+str(i)], 3, len(dataframes['wf'+str(i)]))

wfsdf = pd.concat([sendf['wfsdf2'],sendf['wfsdf3'],sendf['wfsdf4'],sendf['wfsdf5'],sendf['wfsdf6'],sendf['wfsdf7'],
                   sendf['wfsdf8'],sendf['wfsdf9'],sendf['wfsdf10'],sendf['wfsdf11'],sendf['wfsdf12'],
                   sendf['wfsdf13'],sendf['wfsdf14'],sendf['wfsdf15']],axis=1)
wfsdf.columns = range(2,16,1)

for i in range(2,16):
    sendf[f'wwfsdf{i}'] = sens_test(dataframes['wf'+str(i)], 2, len(dataframes['wf'+str(i)]))

wwfsdf = pd.concat([sendf['wwfsdf2'],sendf['wwfsdf3'],sendf['wwfsdf4'],sendf['wwfsdf5'],sendf['wwfsdf6'],sendf['wwfsdf7'],
                   sendf['wwfsdf8'],sendf['wwfsdf9'],sendf['wwfsdf10'],sendf['wwfsdf11'],sendf['wwfsdf12'],
                   sendf['wwfsdf13'],sendf['wwfsdf14'],sendf['wwfsdf15']],axis=1)
wwfsdf.columns = range(2,16,1)

wfsdf = wfsdf.T
wwfsdf = wwfsdf.T





"""

gamma

"""



gamaasendf = {}

for i in range(2,16):
    gamaasendf[f'drsdf{i}'] = sens_test(gammadataframes['dr'+str(i)], 3, len(gammadataframes['dr'+str(i)]))

for i in range(2,16):
    gamaasendf[f'wdrsdf{i}'] = sens_test(gammadataframes['dr'+str(i)], 2, len(gammadataframes['dr'+str(i)]))


gammadrsdf = pd.concat([gamaasendf['drsdf2'], gamaasendf['drsdf3'],gamaasendf['drsdf4'],gamaasendf['drsdf5'],gamaasendf['drsdf6'],gamaasendf['drsdf7'],
                   gamaasendf['drsdf8'],gamaasendf['drsdf9'],gamaasendf['drsdf10'],gamaasendf['drsdf11'],gamaasendf['drsdf12'],
                   gamaasendf['drsdf13'],gamaasendf['drsdf14'],gamaasendf['drsdf15']],axis=1)
gammadrsdf.columns = range(2,16,1)

gammawdrsdf = pd.concat([gamaasendf['wdrsdf2'],gamaasendf['wdrsdf3'],gamaasendf['wdrsdf4'],gamaasendf['wdrsdf5'],gamaasendf['wdrsdf6'],gamaasendf['wdrsdf7'],
                   gamaasendf['wdrsdf8'],gamaasendf['wdrsdf9'],gamaasendf['wdrsdf10'],gamaasendf['wdrsdf11'],gamaasendf['wdrsdf12'],
                   gamaasendf['wdrsdf13'],gamaasendf['wdrsdf14'],gamaasendf['wdrsdf15']],axis=1)
gammawdrsdf.columns = range(2,16,1)

gammadrsdf = gammadrsdf.T
gammawdrsdf = gammawdrsdf.T




for i in range(2,16):
    gamaasendf[f'wfsdf{i}'] = sens_test(gammadataframes['wf'+str(i)]/100, 3, len(gammadataframes['wf'+str(i)]))

gammawfsdf = pd.concat([gamaasendf['wfsdf2'],gamaasendf['wfsdf3'],gamaasendf['wfsdf4'],gamaasendf['wfsdf5'],gamaasendf['wfsdf6'],gamaasendf['wfsdf7'],
                   gamaasendf['wfsdf8'],gamaasendf['wfsdf9'],gamaasendf['wfsdf10'],gamaasendf['wfsdf11'],gamaasendf['wfsdf12'],
                   gamaasendf['wfsdf13'],gamaasendf['wfsdf14'],gamaasendf['wfsdf15']],axis=1)
gammawfsdf.columns = range(2,16,1)


for i in range(2,16):
    gamaasendf[f'wwfsdf{i}'] = sens_test(gammadataframes['wf'+str(i)]/100, 2, len(gammadataframes['wf'+str(i)]))

gammawwfsdf = pd.concat([gamaasendf['wwfsdf2'], gamaasendf['wwfsdf3'],gamaasendf['wwfsdf4'],gamaasendf['wwfsdf5'],gamaasendf['wwfsdf6'],gamaasendf['wwfsdf7'],
                   gamaasendf['wwfsdf8'],gamaasendf['wwfsdf9'],gamaasendf['wwfsdf10'],gamaasendf['wwfsdf11'],gamaasendf['wwfsdf12'],
                   gamaasendf['wwfsdf13'],gamaasendf['wwfsdf14'],gamaasendf['wwfsdf15']],axis=1)
gammawwfsdf.columns = range(2,16,1)

gammawfsdf = gammawfsdf.T
gammawwfsdf = gammawwfsdf.T




"""


tau


"""




tausendf = {}

for i in range(2,16):
    tausendf[f'drsdf{i}'] = sens_test(taudataframes['dr'+str(i)], 3, len(taudataframes['dr'+str(i)]))

for i in range(2,16):
    tausendf[f'wdrsdf{i}'] = sens_test(taudataframes['dr'+str(i)], 2, len(taudataframes['dr'+str(i)]))


taudrsdf = pd.concat([tausendf['drsdf2'],tausendf['drsdf3'],tausendf['drsdf4'],tausendf['drsdf5'],tausendf['drsdf6'],tausendf['drsdf7'],
                   tausendf['drsdf8'],tausendf['drsdf9'],tausendf['drsdf10'],tausendf['drsdf11'],tausendf['drsdf12'],
                   tausendf['drsdf13'],tausendf['drsdf14'],tausendf['drsdf15']],axis=1)
taudrsdf.columns = range(2,16,1)

tauwdrsdf = pd.concat([tausendf['wdrsdf2'],tausendf['wdrsdf3'],tausendf['wdrsdf4'],tausendf['wdrsdf5'],tausendf['wdrsdf6'],tausendf['wdrsdf7'],
                   tausendf['wdrsdf8'],tausendf['wdrsdf9'],tausendf['wdrsdf10'],tausendf['wdrsdf11'],tausendf['wdrsdf12'],
                   tausendf['wdrsdf13'],tausendf['wdrsdf14'],tausendf['wdrsdf15']],axis=1)
tauwdrsdf.columns = range(2,16,1)

taudrsdf = taudrsdf.T
tauwdrsdf = tauwdrsdf.T

for i in range(2,16):
    tausendf[f'wfsdf{i}'] = sens_test(taudataframes['wf'+str(i)]/100, 3, len(taudataframes['wf'+str(i)]))

tauwfsdf = pd.concat([tausendf['wfsdf2'], tausendf['wfsdf3'],tausendf['wfsdf4'],tausendf['wfsdf5'],tausendf['wfsdf6'],tausendf['wfsdf7'],
                   tausendf['wfsdf8'],tausendf['wfsdf9'],tausendf['wfsdf10'],tausendf['wfsdf11'],tausendf['wfsdf12'],
                   tausendf['wfsdf13'],tausendf['wfsdf14'],tausendf['wfsdf15']],axis=1)
tauwfsdf.columns = range(2,16,1)

for i in range(2,16):
    tausendf[f'wwfsdf{i}'] = sens_test(taudataframes['wf'+str(i)]/100, 2, len(taudataframes['wf'+str(i)]))

tauwwfsdf = pd.concat([tausendf['wwfsdf2'],tausendf['wwfsdf3'],tausendf['wwfsdf4'],tausendf['wwfsdf5'],tausendf['wwfsdf6'],tausendf['wwfsdf7'],
                   tausendf['wwfsdf8'],tausendf['wwfsdf9'],tausendf['wwfsdf10'],tausendf['wwfsdf11'],tausendf['wwfsdf12'],
                   tausendf['wwfsdf13'],tausendf['wwfsdf14'],tausendf['wwfsdf15']],axis=1)
tauwwfsdf.columns = range(2,16,1)

tauwfsdf = tauwfsdf.T
tauwwfsdf = tauwwfsdf.T



"""



Drought simulation results

Fig3a


"""


lllist = [x for x in range(2,15)]
llist = [x -0.25 for x in lllist]

error1 = [wdrsdf['vm']-wdrsdf['vb'],wdrsdf['vt']-wdrsdf['vm']]
error2 = [tauwdrsdf['vm']-tauwdrsdf['vb'],tauwdrsdf['vt']-tauwdrsdf['vm']]

plt.figure(figsize=(10,5),dpi=400)
plt.errorbar(lllist,wdrsdf['vm'][lllist],[error1[0][lllist],error1[1][lllist]],color='k',capsize = 5,linestyle='none',marker='o',label='$\\tau(0.068) x Factor,\\gamma(0.034) x 1/Factor$',markersize=11)
plt.errorbar(llist,tauwdrsdf['vm'][lllist],[error2[0][lllist],error2[1][lllist]],color='red',capsize = 5,linestyle='none',marker='^',label='$\\tau(0.068) x Factor$',markersize=11)

plt.xticks(range(2,15,1))
plt.axhline(y=0,color='black',linestyle='-',linewidth=2)
plt.ylim([-15,50])
plt.axhline(y=-10,color='black',linestyle='--',linewidth=1)
plt.axhline(y=10,color='black',linestyle='--',linewidth=1)
plt.axhline(y=20,color='black',linestyle='--',linewidth=1)
plt.axhline(y=30,color='black',linestyle='--',linewidth=1)
plt.axhline(y=40,color='black',linestyle='--',linewidth=1)

plt.xlabel('Factor')
ax = plt.gca()
handles, labels = ax.get_legend_handles_labels()
new_handles = []
for h in handles:
    #only need to edit the errorbar legend entries
    if isinstance(h, container.ErrorbarContainer):
        new_handles.append(h[0])
    else:
        new_handles.append(h)
ax.legend(new_handles, labels,prop = { "size": 15 }, loc ="lower right",framealpha=1,edgecolor='black')
plt.show()

"""


Fig3b


"""

mlist = [2,5,8,11,14][::-1]
mmlist= [22,40,58,76,94][::-1]

plt.figure(figsize=(5,5),dpi=400)

plt.scatter(wdrsdf['vm'][lllist],wdrsdf['d'][lllist],color='k',s=llist)
plt.scatter(tauwdrsdf['vm'][lllist],tauwdrsdf['d'][lllist],color='red',s=llist,marker='^')

for i in range(len(mlist)):
    plt.scatter(100,100,color='k',s=mmlist[i],label=mlist[i])

a, b= np.polyfit(wdrsdf['vm'][lllist],wdrsdf['d'][lllist],1)
# y = -0.37x + 19
plt.plot(range(0,40),a*range(0,40)+b,color='black',linestyle='--')
a, b= np.polyfit(tauwdrsdf['vm'][lllist],tauwdrsdf['d'][lllist],1)
# y = -0.32x + 17

plt.plot(range(0,40),a*range(0,40)+b,color='red',linestyle='--')
plt.legend(prop={"size": 9}, loc="upper right", framealpha=1, edgecolor='black')
plt.ylim([0,35])
plt.xlim([0,40])
plt.show()


"""



Wildfire simulation results
Fig3c


"""


lllist = [x for x in range(2,15)]
llist = [x -0.25 for x in lllist]

error1 = [wwfsdf['vm']-wwfsdf['vb'],wwfsdf['vt']-wwfsdf['vm']]
error2 = [tauwwfsdf['vm']-tauwwfsdf['vb'],tauwwfsdf['vt']-tauwwfsdf['vm']]

plt.figure(figsize=(10,5),dpi=400)
plt.errorbar(lllist,wwfsdf['vm'][lllist],[error1[0][lllist],error1[1][lllist]],color='k',capsize = 5,linestyle='none',marker='o',
             label='$\\tau(0.046) x Factor,\\gamma(0.026) x 1/Factor$',markersize=11,)
plt.errorbar(llist,tauwwfsdf['vm'][lllist],[error2[0][lllist],error2[1][lllist]],color='red',capsize = 5,linestyle='none',marker='^',
             label='$\\tau(0.046) x Factor$',markersize=11)

plt.xticks(range(2,15,1))
plt.axhline(y=0,color='black',linestyle='-',linewidth=2)
plt.ylim([-15,50])
plt.xlabel('Factor')
plt.axhline(y=-10,color='black',linestyle='--',linewidth=1)
plt.axhline(y=10,color='black',linestyle='--',linewidth=1)
plt.axhline(y=20,color='black',linestyle='--',linewidth=1)
plt.axhline(y=30,color='black',linestyle='--',linewidth=1)
plt.axhline(y=40,color='black',linestyle='--',linewidth=1)
ax = plt.gca()
handles, labels = ax.get_legend_handles_labels()
new_handles = []
for h in handles:
    #only need to edit the errorbar legend entries
    if isinstance(h, container.ErrorbarContainer):
        new_handles.append(h[0])
    else:
        new_handles.append(h)
ax.legend(new_handles, labels,prop = { "size": 15 }, loc ="lower right",framealpha=1,edgecolor='black')
plt.show()



"""

Fig3d

"""



plt.figure(figsize=(5,5),dpi=400)

plt.scatter(wwfsdf['vm'][lllist],wwfsdf['d'][lllist],color='k',s=llist)
plt.scatter(tauwwfsdf['vm'][lllist],tauwwfsdf['d'][lllist],color='red',s=llist,marker='^')
for i in range(len(mlist)):
    plt.scatter(100,100,color='k',s=mmlist[i],label=mlist[i])

a, b= np.polyfit(wwfsdf['vm'][lllist],wwfsdf['d'][lllist],1)
# y = -0.73x + 37
plt.plot(range(0,40),a*range(0,40)+b,color='black',linestyle='--')
a, b= np.polyfit(tauwwfsdf['vm'][lllist],tauwwfsdf['d'][lllist],1)
# y = -0.58 + 32

plt.plot(range(0,40),a*range(0,40)+b,color='red',linestyle='--')
plt.legend(prop={"size": 9}, loc="upper right", framealpha=1, edgecolor='black')

plt.ylim([0,35])
plt.xlim([0,40])
plt.show()
